#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/ec.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#include "sm2_crypto.h"

/* Function: sm2_keypair_generate
 * Description: Generates SM2 public and private key
 * Parameters:
 * ppKeypair - pointer to key pair
 * Returns: 0 success, -1 failure
 */
int sm2_keypair_generate(EVP_PKEY **ppKeypair)
{
    EVP_PKEY_CTX *pPkeyCtx = NULL;
    
    /* Check key pair pointer */
    if (ppKeypair == NULL) {
        fprintf(stderr, "sm2_keypair_generate: key pair buffer is null.\n");
        return (-1);
    }

    /* Create key pair context */
    pPkeyCtx = EVP_PKEY_CTX_new_id(EVP_PKEY_SM2, NULL);
    if (pPkeyCtx == NULL) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize key generation */
    if (EVP_PKEY_keygen_init(pPkeyCtx) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Set SM2 curve parameters */
    if (EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pPkeyCtx, NID_sm2) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Generate key pair */
    if (EVP_PKEY_keygen(pPkeyCtx, ppKeypair) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Clean context */
    EVP_PKEY_CTX_free(pPkeyCtx);
    return (0);
}

/* Function: sm2_encrypt
 * Description: SM2 encrypt using SM2 public key
 * Parameters:
 * pPublicKey       - pointer to SM2 public key buffer
 * pucID            - pointer to SM2 ID
 * stIdLen          - SM2 ID length
 * pucPlaintext     - pointer to plaintext buffer
 * stPlaintextLen   - plaintext length
 * pucCiphertext   -  pointer to ciphertext buffer
 * pstCiphertextLen - pointer to ciphertext length
 * Returns: 0 success, -1 failure
 */
int sm2_encrypt(EVP_PKEY *pPublicKey,
                const unsigned char *pucId,
                size_t stIdLen,
                const unsigned char *pucPlaintext,
                size_t stPlaintextLen,
                unsigned char *pucCiphertext,
                size_t *pstCiphertextLen)
{
    EVP_PKEY_CTX *pPkeyCtx = NULL;
    size_t stCiphertextLen = 0;

    /* Check public key buffer */
    if (pPublicKey == NULL) {
        fprintf(stderr, "sm2_encrypt: public key buffer is null.\n");
        return (-1);
    }

    /* Check ID buffer */
    if (pucId == NULL) {
        fprintf(stderr, "sm2_encrypt:: ID buffer is null.\n");
        return (-1);
    }

    /* Check ID length */
    if (stIdLen == 0) {
        fprintf(stderr, "sm2_encrypt: invalid ID length %zu.\n", stIdLen);
        return (-1);
    }

    /* Check plaintext buffer */
    if (pucPlaintext == NULL) {
        fprintf(stderr, "sm2_encrypt:: plaintext buffer is null.\n");
        return (-1);
    }

    /* Check plaintext length */
    if (stPlaintextLen == 0) {
        fprintf(stderr, "sm2_encrypt: invalid plaintext length %zu.\n", stPlaintextLen);
        return (-1);
    }

    /* Check ciphertext buffer */
    if (pucCiphertext == NULL) {
        fprintf(stderr, "sm2_encrypt: ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length buffer */
    if (pstCiphertextLen == NULL) {
        fprintf(stderr, "sm2_encrypt: ciphertext length buffer is null.\n");
        return (-1);
    }

    /* Create key pair context */
    pPkeyCtx = EVP_PKEY_CTX_new(pPublicKey, NULL);
    if (pPkeyCtx == NULL) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Set ID */
    if (EVP_PKEY_CTX_set1_id(pPkeyCtx, pucId, stIdLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Initialize key pair encrypt */
    if (EVP_PKEY_encrypt_init(pPkeyCtx) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Process plaintext */
    if (EVP_PKEY_encrypt(pPkeyCtx, NULL, &stCiphertextLen, pucPlaintext, stPlaintextLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Check ciphertext buffer length */
    if ((*pstCiphertextLen) < stCiphertextLen) {
        fprintf(stderr, "sm2_encrypt: ciphertext buffer length %zu is too small.\n", (*pstCiphertextLen));
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Encryp */
    if (EVP_PKEY_encrypt(pPkeyCtx, pucCiphertext, &stCiphertextLen, pucPlaintext, stPlaintextLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Set ciphertext length */
    *pstCiphertextLen = stCiphertextLen;
    
    /* Clean context */
    EVP_PKEY_CTX_free(pPkeyCtx);
    return (0);
}

/* Function: sm2_decrypt
 * Description: SM2 decrypts using SM2 private key
 * Parameters:
 * pPrivateKey     - pointer to SM2 private key buffer
 * pucID           - pointer to SM2 ID
 * stIdLen         - SM2 ID length
 * pucCiphertext   - pointer to ciphertext buffer
 * stCiphertextLen - ciphertext length
 * pucPlaintext   -  pointer to plaintext buffer
 * pstPlaintextLen - pointer to plaintext length
 * Returns: 0 success, -1 failure
 */
int sm2_decrypt(EVP_PKEY *pPrivateKey,
                const unsigned char *pucId,
                size_t stIdLen,
                const unsigned char *pucCiphertext,
                size_t stCiphertextLen,
                unsigned char *pucPlaintext,
                size_t *pstPlaintextLen)
{
    EVP_PKEY_CTX *pPkeyCtx = NULL;
    size_t stPlaintextLen = 0;

    /* Check private key buffer */
    if (pPrivateKey == NULL) {
        fprintf(stderr, "sm2_decrypt：private key buffer is null.\n");
        return (-1);
    }

    /* Check ID buffer */
    if (pucId == NULL) {
        fprintf(stderr, "sm2_encrypt:: ID buffer is null.\n");
        return (-1);
    }

    /* Check ID length */
    if (stIdLen == 0) {
        fprintf(stderr, "sm2_encrypt: invalid ID length %zu.\n", stIdLen);
        return (-1);
    }

    /* Check ciphertext buffer */
    if (pucCiphertext == NULL) {
        fprintf(stderr, "sm2_decrypt：ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length */
    if (stCiphertextLen == 0) {
        fprintf(stderr, "sm2_decrypt：invalid ciphertext length %zu.\n", stCiphertextLen);
        return (-1);
    }

    /* Check plaintext buffer */
    if (pucPlaintext == NULL) {
        fprintf(stderr, "sm2_decrypt：plaintext buffer is null.\n");
        return (-1);
    }

    /* Check plaintext length buffer */
    if (pstPlaintextLen == NULL) {
        fprintf(stderr, "sm2_decrypt：plaintext length buffer is null.\n");
        return (-1);
    }

    /* Create key pair context */
    pPkeyCtx = EVP_PKEY_CTX_new(pPrivateKey, NULL);
    if (pPkeyCtx == NULL) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Set ID */
    if (EVP_PKEY_CTX_set1_id(pPkeyCtx, pucId, stIdLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Initialize key pair decrypt */
    if (EVP_PKEY_decrypt_init(pPkeyCtx) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Process plaintext */
    if (EVP_PKEY_decrypt(pPkeyCtx, NULL, &stPlaintextLen, pucCiphertext, stCiphertextLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Check plaintext buffer length */
    if ((*pstPlaintextLen) < stPlaintextLen) {
        fprintf(stderr, "sm2_decrypt: plaintext buffer length %zu is too small.\n", (*pstPlaintextLen));
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Decrypt */
    if (EVP_PKEY_decrypt(pPkeyCtx, pucPlaintext, &stPlaintextLen, pucCiphertext, stCiphertextLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_PKEY_CTX_free(pPkeyCtx);
        return (-1);
    }

    /* Set plaintext length */
    *pstPlaintextLen = stPlaintextLen;
    
    /* Clean context */
    EVP_PKEY_CTX_free(pPkeyCtx);
    return (0);
}

/* Function: sm2_sign
 * Description: Generates SM2 signature using private key
 * Parameters:
 * pPrivateKey - pointer to SM2 private key buffer
 * pucID       - pointer to SM2 ID
 * stIdLen     - SM2 ID length
 * pucMsg      - pointer to message
 * stMsgLen    - message length
 * pucSign     - pointer to signature buffer
 * pstSignLen  - pointer to signature length
 * Returns: 0 success, -1 failure
 */
int sm2_sign(EVP_PKEY *pPrivateKey,
            const unsigned char *pucId,
            size_t stIdLen,
            const unsigned char *pucMsg,
            size_t stMsgLen,
            unsigned char *pucSign,
            size_t *pstSignLen)
{
    EVP_MD_CTX *pMdCtx = NULL;
    EVP_PKEY_CTX *pPkeyCtx = NULL;
    size_t stSignLen = 0;

    /* Check private key buffer */
    if (pPrivateKey == NULL) {
        fprintf(stderr, "sm2_sign：private key buffer is null.\n");
        return (-1);
    }

    /* Check ID buffer */
    if (pucId == NULL) {
        fprintf(stderr, "sm2_encrypt:: ID buffer is null.\n");
        return (-1);
    }

    /* Check ID length */
    if (stIdLen == 0) {
        fprintf(stderr, "sm2_encrypt: invalid ID length %zu.\n", stIdLen);
        return (-1);
    }

    /* Check message buffer */
    if (!pucMsg) {
        fprintf(stderr, "sm2_sign: message buffer is null.\n");
        return (-1);
    }

    /* Check message length */
    if (stMsgLen == 0) {
        fprintf(stderr, "sm2_sign: invalid message length %zu.\n", stMsgLen);
        return (-1);
    }

    /* Check signature buffer */
    if (pucSign == NULL) {
        fprintf(stderr, "sm2_sign：signature buffer is null.\n");
        return (-1);
    }

    /* Check signature length */
    if (pstSignLen == NULL) {
        fprintf(stderr, "sm2_sign：signature length buffer is null.\n");
        return (-1);
    }

    /* Create MD5 context */
    pMdCtx = EVP_MD_CTX_new();
    if (pMdCtx == NULL) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize digest signature with SM3 */
    if (EVP_DigestSignInit(pMdCtx, NULL, EVP_sm3(), NULL, pPrivateKey) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Get key pair context */
    pPkeyCtx = EVP_MD_CTX_get_pkey_ctx(pMdCtx);
    if (!pPkeyCtx) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Set ID */
    if (EVP_PKEY_CTX_set1_id(pPkeyCtx, pucId, stIdLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Process message */
    if (EVP_DigestSign(pMdCtx, NULL, &stSignLen, pucMsg, stMsgLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Check signature buffer length */
    if ((*pstSignLen) < stSignLen) {
        fprintf(stderr, "sm2_sign: signature buffer length %zu is too small.\n", (*pstSignLen));
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Generate actual signature */
    if (EVP_DigestSign(pMdCtx, pucSign, &stSignLen, pucMsg, stMsgLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Set signature length */
    *pstSignLen = stSignLen;
    
    /* Clean context */
    EVP_MD_CTX_free(pMdCtx);

    return (0);
}

/* Function: sm2_verify
 * Description: Verifies SM2 signature using public key
 * Parameters:
 * pPublicKey - pointer to SM2 public key buffer
 * pucID      - pointer to SM2 ID
 * stIdLen    - SM2 ID length
 * pucMsg     - pointer to message buffer
 * stMsgLen   - message length
 * pucSign    - pointer to signature buffer
 * stSignLen  - Signature length
 * Returns: 0 invalid, 1 valid, -1 failure
 */
int sm2_verify(EVP_PKEY *pPublicKey,
               const unsigned char *pucId,
               size_t stIdLen,
               const unsigned char *pucMsg,
               size_t stMsgLen,
               const unsigned char *pucSign,
               size_t stSignLen)
{
    EVP_MD_CTX *pMdCtx = NULL;
    EVP_PKEY_CTX *pPkeyCtx = NULL;
    int iRet = 0;

    /* Check public key buffer */
    if (pPublicKey == NULL) {
        fprintf(stderr, "sm2_verify：public key buffer is null.\n");
        return (-1);
    }

    /* Check ID buffer */
    if (pucId == NULL) {
        fprintf(stderr, "sm2_encrypt:: ID buffer is null.\n");
        return (-1);
    }

    /* Check ID length */
    if (stIdLen == 0) {
        fprintf(stderr, "sm2_encrypt: invalid ID length %zu.\n", stIdLen);
        return (-1);
    }

    /* Check message buffer */
    if (!pucMsg) {
        fprintf(stderr, "sm2_verify: message buffer is null.\n");
        return (-1);
    }

    /* Check message length */
    if (stMsgLen == 0) {
        fprintf(stderr, "sm2_verify: invalid message length %zu.\n", stMsgLen);
        return (-1);
    }

    /* Check signature buffer */
    if (pucSign == NULL) {
        fprintf(stderr, "sm2_verify：signature buffer is null.\n");
        return (-1);
    }

    /* Check signature length */
    if (stSignLen == 0) {
        fprintf(stderr, "sm2_verify：invalid signature length %zu.\n", stSignLen);
        return (-1);
    }

    /* Create MD5 context */
    pMdCtx = EVP_MD_CTX_new();
    if (pMdCtx == NULL) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize digest verification with SM3 hash */
    if (EVP_DigestVerifyInit(pMdCtx, NULL, EVP_sm3(), NULL, pPublicKey) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Get key pair context */
    pPkeyCtx = EVP_MD_CTX_get_pkey_ctx(pMdCtx);
    if (!pPkeyCtx) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Set ID */
    if (EVP_PKEY_CTX_set1_id(pPkeyCtx, pucId, stIdLen) <= 0) {
        ERR_print_errors_fp(stderr);
        EVP_MD_CTX_free(pMdCtx);
        return (-1);
    }

    /* Verify */
    iRet = EVP_DigestVerify(pMdCtx, pucSign, stSignLen, pucMsg, stMsgLen);
    
    /* Clean context */
    EVP_MD_CTX_free(pMdCtx);

    /* Check result */
    if (iRet == 0) {
        /* Signature is invalid */
        fprintf(stderr, "sm2_verify: signature verification failed.\n");
        return (0);
    } else if (iRet == 1) {
        /* Signature is valid */
        return (1);
    } else {
        /* Error. */
        ERR_print_errors_fp(stderr);
        return (-1);
    }
}